import sys
sys.path.append("..")
from Dataloader.twitterloader import TwitterSet, SentiReader
import torch.nn as nn
from Dataloader.dataloader_utils import Sort_data
from SentModel.Sent2Vec import W2VRDMVec
from PropModel.SeqPropagation import GRUModel
from torch.utils.data import DataLoader
from RumdetecFramework.ReinforceRumorFramework import ReinforceRumorDetection, SeqRumorDetecEnv

tr_set = TwitterSet()
tr_set.load_data_fast(data_prefix="../../data/twitter_tr2", min_len=5)

dev = TwitterSet()
dev.load_data_fast(data_prefix="../../data/twitter_dev2", min_len=5)

te = TwitterSet()
te.load_data_fast(data_prefix="../../data/twitter_te2", min_len=5)

tr_set, dev, te = Sort_data(tr_set, dev, te)

sent2vec1 = W2VRDMVec("../../saved/glove_en/", 300)
sent2vec2 = W2VRDMVec("../../saved/glove_en/", 300)
prop = GRUModel(300, 256, 1, 0.2)

actor = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 2)
        )
critic = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

subrrd = ReinforceRumorDetection(sent2vec1, sent2vec2, prop, actor, critic)
env = SeqRumorDetecEnv(tr_set, num_worker=20,  hidden_size=prop.prop_hidden_size)
dev_loader = DataLoader(dev, batch_size=20, shuffle=False, collate_fn=dev.collate_raw_batch)
te_loader = DataLoader(te, batch_size=20, shuffle=False, collate_fn=te.collate_raw_batch)

subj_tr = SentiReader("../../data/sub_train.csv")
subj_dev = SentiReader("../../data/sub_dev.csv")
subj_te = SentiReader("../../data/sub_test.csv")

subrrd.joint_train_iters(tr_set, dev, te, subj_tr, subj_dev, subj_te,
                    valid_every=100, max_epochs=50, lr_discount=1.0,
                    best_valid_acc=0.0, best_test_acc=0.0, best_valid_test_acc=0.0,
                    log_dir="../../logs", log_suffix="_RumorDetection", model_file="../../saved/SubRDM_Sort.pkl")
# subrrd.load_model("../../saved/SubRDM.pkl")
# subrrd.valid(te_loader, all_metrics=True)
# subrrd.ACTrain(env, dev_loader, te_loader, True, 10000000)